import numpy as np
import os
import output_function as of
import matplotlib.pyplot as plt

def SOL_gamma_n(n_density,psi):
    """
    相当于公式中的\Gamma_n,体现SOL效应,在n方程中
    phi=-psi
    alpha_sh0=Omega_freq/L_paral
    """
    c1=L_perp**2/L_paral
    n_SOL_effect=np.zeros_like(n_density)
    for j in range(1,Ny-1):
        for i in range(1,Nx-1):
            n_SOL_effect[j][i]=h_array[j]*(-alpha_sh0*(n_density[j][i]-n_0)-c1*n_density[j][i]*psi[j][i])
    return n_SOL_effect[1:-1, 1:-1]

def convection_operator(myfun,u,v,dx,dy):
    """
    计算对流算子。
    :param myfun: 二维数组，要计算对流算子的函数。
    :param u: 二维数组，x方向的速度。
    :param v: 二维数组，y方向的速度。
    :param dx: x方向的网格大小。
    :param dy: y方向的网格大小。
    :return: 对流算子的结果。
    """
    """x向对流速度u造成的输运,u \partial_x f"""
    uw_x = 0.5*u[1:-1,1:-1]  * (myfun[1:-1, 2:] - myfun[1:-1, :-2])/ dx
    """v \partial_y f"""
    vw_y = 0.5*v[1:-1,1:-1]  * (myfun[2:,1:-1] - myfun[:-2,1:-1])/ dy

    #守恒形式
    # uw_x =  (u[1:-1, 2:]*myfun[1:-1, 2:] -u[1:-1, :-2]*myfun[1:-1, :-2])/ 2 / dx
    # vw_y =  (v[2:, 1:-1]*myfun[2:, 1:-1] - v[:-2, 1:-1]*myfun[:-2, 1:-1])/2 / dy

    return uw_x,vw_y

# 使用示例
base_path = './data/SOL_Ra=75000.0_Lx_5.0_Dn=0.04_mu=0.1/'  # 根据需要修改为您的保存路径
out_iter = 19000  # 假设我们要加载第300时间步的数据

base_path = './data/SOL_Ra=60000.0_Lx_5.0_Dn=0.1_mu=0.05/'  # 根据需要修改为您的保存路径
out_iter = 15000  # 假设我们要加载第300时间步的数据

# base_path = './data/SOL_Ra=30000.0_Lx_5.0_Dn=0.1_mu=0.1/'  # 根据需要修改为您的保存路径
# out_iter = 20000  # 假设我们要加载第300时间步的数据

base_path = './data/SOL_Ra=7500.0_Lx_5.0_Dn=0.2_mu=0.2/'  # 根据需要修改为您的保存路径
out_iter = 44000  # 假设我们要加载第300时间步的数据

# base_path = './data/SOL_Ra=120000.0_Lx_5.0_Dn=0.05_mu=0.05/'  # 根据需要修改为您的保存路径
# out_iter = 19000  # 假设我们要加载第300时间步的数据

# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')

n_density=n

# 加载模拟参数
config = of.load_simulation_config(f'{base_path}simulation_config.json')
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Ly_in=config['Ly_in']
Ly_out=config['Ly_out']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
L_perp = config['L_perp']
L_paral = config['L_paral']
g_hat = config['g_hat']
zeta = config['zeta']
D_n = config['D_n']
mu = config['mu']
Ra_star = config['Ra_star']
Pr = config['Pr']
Omega_freq = config['Omega_freq']
alpha_sh0 = config['alpha_sh0']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

# 打印一些变量以确认它们已被正确加载和赋值
print(f"dx: {dx}")
print(f"Ra_star: {Ra_star}")
print(f"dt: {dt}")
# 确认是否成功加载了正确的时间步长
if out_iter_loaded == out_iter:
    print(f"Successfully loaded data for time step: {out_iter_loaded}")
    print(y)
else:
    print("Failed to load the correct time step or file does not exist.")

# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
h_array=np.zeros_like(y)
for j in range(0,Ny):
    h_array[j]=np.heaviside((y[j]-Ly_in),0.5)
#of.diag_n_pcolor(n,n_init0_matrix,u,v,xx,yy,out_iter,Ra_star,dt)
#of.diag_n_origin_pcolor_filepath(n,u,v,xx,yy,out_iter,Ra_star,dt)
t_total=out_iter*dt #总时间

# 初始化n_ave数组，用于存储每个j值对应的n的平均值
n_ave = np.zeros(Ny-1)  # Ny为n数组在y轴方向的长度
n_init_ave=np.zeros(Ny-1)
convec_n=np.zeros(Ny-1)
convec_n2=np.zeros(Ny-1)
diff_n=np.zeros(Ny-1)
total_flux=np.zeros(Ny-1)
radial_x=np.zeros(Ny-1)
interchange_loss=np.zeros(Ny-1)
SOL_loss=np.zeros(Ny-1)

n_SOL_effect=np.zeros_like(n_density)
n_SOL_effect[1:-1,1:-1]=SOL_gamma_n(n_density,psi)

un_x,vn_y=convection_operator(n,u,v,dx,dy)
# 对于每个j值，计算对应列的平均值
for j in range(Ny-1):
    n_ave[j] = np.mean(n[j, :])
    n_init_ave[j]=np.mean(n_init0_matrix[j, :])
    radial_x[j]=y[j]-Ly_in
    if j>0:
        #需要积分的项
        interchange_loss[j]=interchange_loss[j-1]
        SOL_loss[j]=SOL_loss[j-1]
        convec_n2[j]=convec_n2[j-1]
        convec_n2[j]+=np.mean(vn_y[j-1,:])*Lx*dy  #直接积分
    for i in range(Nx):
        convec_n[j]+=n[j,i]*v[j,i]*dx   #采用nv直接计算
        diff_n[j]+=(n[j,i]-n[j+1,i])/dy*dx
        #需要积分的项
        interchange_loss[j]+=zeta*n[j,i]*v[j,i]*dx*dy
        SOL_loss[j]+=-n_SOL_effect[j,i]*dx*dy
    #需要积分的
    total_flux[j]=convec_n[j]+diff_n[j]+interchange_loss[j]+SOL_loss[j]

# 绘制n_ave与y的关系图
plt.figure(figsize=(10, 6))
plt.plot(radial_x, n_ave, label='Convective Density profile', color='red')
plt.plot(radial_x, n_init_ave, label='Diffusive Density profile', color='blue')
plt.axvline(x=0, color='black', linestyle='--', linewidth=1, label='LCFS')
plt.xlabel('radial position relative to LCFS')
plt.ylabel('Average Density ')
plt.title(f'Convective vs. Diffusive Density profile,Ra={Ra_star},t={t_total:.3f}')
plt.legend()
plt.grid(True)
plt.show()

# 绘制n_ave与y的关系图
plt.figure(figsize=(8, 4))
plt.plot(radial_x, convec_n, label='<nv>', color='red')
# plt.plot(radial_x, convec_n2, label='<nv>2', color='y')
plt.plot(radial_x, diff_n, label='-<dn/dx>', color='blue')
# plt.plot(radial_x, total_flux, label='total_flux', color='purple')
# plt.plot(radial_x, interchange_loss, label='interchange_loss', color='green')
# plt.plot(radial_x, SOL_loss, label='SOL_loss', color='pink')
plt.axvline(x=0, color='black', linestyle='--', linewidth=1, label='LCFS')
plt.xlim(-0.4,1)
plt.xlabel('radial position relative to LCFS')
plt.ylabel('Average Density ')
plt.title(f'Convective vs. Diffusive transport,Ra={Ra_star},t={t_total:.3f}')
plt.legend(loc='right')
plt.grid(True)
plt.show()

exit(0)
yy1=yy-1
flag_pcolormesh_stream=True
if flag_pcolormesh_stream:
    t_diag=out_iter*dt
    fig, ax = plt.subplots(figsize=(10, 5))
    n_difference=n_density
    # 绘制伪彩色图
    plt.streamplot(xx, yy1, u, v, linewidth=1)
    
    # plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), vmin=-np.abs(n_difference).max(), vmax=np.abs(n_difference).max(), shading='auto')
    # plt.colorbar(label='n',boundaries=np.linspace(-np.abs(n_difference).max(), np.abs(n_difference).max()))
    
    # 使用自动缩放的颜色条
    pcm = plt.pcolormesh(xx, yy1, n_difference, cmap=plt.get_cmap('Spectral_r'), shading='auto')
    # 添加标题
    ax.set_title('Ra={}, t={:.3f}'.format(Ra_star, t_diag), fontsize=16)
    # 添加颜色条
    #cbar = fig.colorbar(cs, orientation='vertical')
    # 设置坐标轴的比例相同，使得图形保持矩形
    ax.set_aspect('equal')
    # 添加颜色条，不需要指定boundaries
    cbar = fig.colorbar(pcm, ax=ax, orientation='vertical' ,fraction=0.02)
    cbar.set_label('n')
    plt.show()

if flag_pcolormesh_stream:
    t_diag=out_iter*dt
    fig, ax = plt.subplots(figsize=(10, 5))
    n_difference=n_density
    # 绘制伪彩色图
    
    #plt.streamplot(xx, yy, u, v, linewidth=1)
    
    # plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), vmin=-np.abs(n_difference).max(), vmax=np.abs(n_difference).max(), shading='auto')
    # plt.colorbar(label='n',boundaries=np.linspace(-np.abs(n_difference).max(), np.abs(n_difference).max()))
    
    # 使用自动缩放的颜色条
    pcm = plt.pcolormesh(xx, yy, v, cmap=plt.get_cmap('Spectral_r'), shading='auto')
    # 添加标题
    ax.set_title('Ra={}, t={}'.format(Ra_star, t_diag), fontsize=16)
    # 添加颜色条
    #cbar = fig.colorbar(cs, orientation='vertical')
    # 设置坐标轴的比例相同，使得图形保持矩形
    ax.set_aspect('equal')
    # 添加颜色条，不需要指定boundaries
    cbar = fig.colorbar(pcm, ax=ax, orientation='vertical' ,fraction=0.02)
    cbar.set_label('v')
    plt.show()

#绘制流函数二维颜色图。
flag_contourf=True
if flag_contourf:
    fig, ax = plt.subplots(figsize=(10, 5))
    # 绘制等值线图，增加levels参数来使等值线更密集  
    levels = np.linspace(psi.min(), psi.max(), 20)  # 这里设置为50个级别  
    cs = ax.contourf(xx, yy1, psi, levels=levels, cmap=plt.get_cmap('Spectral').reversed())  
    plt.streamplot(xx, yy1, u, v, linewidth=0.5,color='black')
    # 添加标题  
    ax.set_title('Ra={},t={:.3f}'.format(Ra_star,t_total), fontsize=16)  # 设置标题文本和字体大小  
    # ax.set_xlabel('X')  # 自行替换为合适的标签
    # ax.set_ylabel('Y')  # 自行替换为合适的标签
    #添加colorbar
    cbar = fig.colorbar(cs,ax=ax, orientation='vertical',fraction=0.02)
    cbar.set_label('psi') 
    ax.set_aspect('equal')
    plt.show()



# fig, ax = plt.subplots()
# q = plt.quiver(xx, yy, u, v, scale=0.002)  # scale factor of 2 to make arrows larger
# ax.quiverkey(q, X=0.3, Y=1.1, U=1000,
#              label='Quiver key, length = 10', labelpos='E')
# plt.show()